package com.tvunetworks.center.device.config.datasource;

import org.apache.commons.lang.StringUtils;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author Ethan
 * @data: 2021/6/15 14:12
 * @program:
 * @description:
 */
public class RoutingDataSource extends AbstractRoutingDataSource {

	private final int dataSourceNumber;
	private AtomicInteger count = new AtomicInteger(0);

	public RoutingDataSource(int dataSourceNumber) {
		this.dataSourceNumber = dataSourceNumber;
	}

	@Override
	protected Object determineCurrentLookupKey() {
		String typeKey = DataSourceContextHolder.getJdbcType();
		if (StringUtils.isBlank(typeKey) || typeKey.equals(DataSourceType.write.getType())) {
			return DataSourceType.write.getType();
		}
		// 读 简单负载均衡
		int number = count.getAndAdd(1);
		int lookupKey = number % dataSourceNumber;
		return  Integer.valueOf(lookupKey);
	}
}

